#include <linux/cdev.h>
#include <linux/module.h>
#include <linux/fs.h>
#include <linux/init.h>
#include <linux/slab.h>
#include <linux/uaccess.h>
#include <linux/errno.h>
#include <linux/rbtree.h>
#include <linux/list.h>


#define CHAR_DRIVER_TEST_DEBUG
#ifdef CHAR_DRIVER_TEST_DEBUG
    #define test_driver_debug(fmt, ...) printk(KERN_DEBUG fmt, ##__VA_ARGS__)
#else
    #define test_driver_debug(fmt, ...)
#endif

enum {
	TEST_CHAR_NORMAL_MODE = 0,
	TEST_CHAR_PIPE_MODE,
	TEST_CHAR_PACKET_MODE,
	TEST_CHAR_ROUND_ROBIN_PACKET_MODE,
};


struct test_char_drv_buf_info {
	unsigned long buf_size;          /* buffer 空间大小 */
	unsigned char *buf_data;         /* buffer  */ 

	unsigned int write_index;        /* 待写地址 */
	unsigned int read_index;         /* 待读地址 */
	unsigned char have_writed;
	
	size_t total_read_len;
	size_t total_write_len;

	struct semaphore sem;            /* 保护信号量 */

	wait_queue_head_t wait_queue;    /* 等待队列 */
};


struct test_char_drv_packet {
	struct list_head list;
	unsigned int len;
	unsigned char *data;
};

struct test_char_drv_packet_info {
	struct test_char_drv_packet packet;
	unsigned int packet_cnt;         /* packet个数 */
	struct mutex packet_mutex;       /* 保护packet链表 */
	struct semaphore sem;            /* 保护信号量 */
	wait_queue_head_t wait_queue;    /* 等待队列 */
};


struct test_char_drv_dev {
#if 0
	unsigned long buf_size;          /* buffer 空间大小 */
	unsigned char *buf_data;         /* buffer  */ 

	unsigned int write_index;        /* 待写地址 */
	unsigned int read_index;         /* 待读地址 */
	unsigned char have_writed;
	
	size_t total_read_len;
	size_t total_write_len;

	struct semaphore sem;

	wait_queue_head_t wait_queue; 

#endif

	struct test_char_drv_buf_info buf_info;
	struct test_char_drv_packet_info packet_info;

	struct kobject *kobj;

	struct radix_tree_root buf_tree;
	spinlock_t tree_lock;
	atomic64_t count_in_buf_tree;

	/* round robin 链表 */
	struct list_head round_head_list;
	struct mutex round_head_list_mutex;
	unsigned int round_index;
	unsigned int round_cnt;
};

struct test_char_priv {		
	struct test_char_drv_dev *dev;
	unsigned mode;

	/* 链入 round robin 链表 */
	struct list_head list;
	struct test_char_drv_packet_info packet_info;
};

static int major_id = 100; /* 主设备号 */
static struct test_char_drv_dev char_dev; 


static int test_char_drv_open(struct inode *inode, struct file *file)
{
	struct test_char_priv *priv;
	//struct task_struct

	test_driver_debug("open inode=0x%x, file=0x%x pid=%d ppid=%d...\n", inode, file, 
		current->pid, current->real_parent->pid);

	priv = kmalloc(sizeof(struct test_char_priv), GFP_KERNEL);
	if (!priv) {
		test_driver_debug("open failed...\n");
		return -1;
	}

	priv->mode = TEST_CHAR_NORMAL_MODE;
	priv->dev  = &char_dev;
	file->private_data = priv;

	test_driver_debug("open success packet_info=0x%x...\n", &priv->packet_info);
	
    return 0;
}


static int test_char_drv_release(struct inode *inode, struct file *file)
{
	unsigned long flags;
	struct test_char_priv *priv = file->private_data;
	struct test_char_drv_dev *dev = priv->dev;
	struct test_char_drv_buf_info *buf_info;
	struct test_char_drv_packet_info *packet_info;
	
	test_driver_debug("release inode=0x%x, file=0x%x pid=%d ppid=%d...\n", inode, file, 
		current->pid, current->real_parent->pid);

	if (priv->mode == TEST_CHAR_PIPE_MODE) {
		buf_info = radix_tree_lookup(&dev->buf_tree, (unsigned long)file);
		if (!buf_info) {
			test_driver_debug("release radix_tree_lookup error file=0x%x...\n", file);
			return -1;
		}

		spin_lock_irqsave(&dev->tree_lock, flags);
		buf_info = radix_tree_delete(&dev->buf_tree, (unsigned long)file);
		atomic64_dec(&dev->count_in_buf_tree);
		spin_unlock_irqrestore(&dev->tree_lock, flags);

		if (buf_info) {
			vfree(buf_info->buf_data);
			kfree(buf_info);
		}
	}
	
	if(priv->mode == TEST_CHAR_ROUND_ROBIN_PACKET_MODE) {
		packet_info = &priv->packet_info;
		
		mutex_lock(&dev->round_head_list_mutex);
		list_del(&(packet_info->packet.list));
		if(dev->round_cnt)
			dev->round_cnt--;
		mutex_unlock(&dev->round_head_list_mutex);
	}

	kfree(priv);
	
	test_driver_debug("release success ...\n");
	
    return 0;
}


static ssize_t test_char_drv_read(struct file *file, char __user *buf,
		    size_t nbytes, loff_t *ppos)
{
	size_t count;
	struct test_char_priv *priv = file->private_data;
	struct test_char_drv_dev *dev = priv->dev;
	struct test_char_drv_buf_info *buf_info; 
	struct test_char_drv_packet_info *packet_info; 
	struct test_char_drv_packet *packet = NULL;

	test_driver_debug("read file=0x%x pid=%d, mode=%d\n", file, current->pid, priv->mode);

	switch (priv->mode)
	{
	case TEST_CHAR_PACKET_MODE:
	case TEST_CHAR_ROUND_ROBIN_PACKET_MODE:
		if (priv->mode == TEST_CHAR_PACKET_MODE) 
			packet_info = &dev->packet_info;
		else
			packet_info = &priv->packet_info;

		mutex_lock(&packet_info->packet_mutex);
		while (!packet_info->packet_cnt) {
			mutex_unlock(&packet_info->packet_mutex);
			if (file->f_flags & O_NONBLOCK)
				return -EAGAIN;

			test_driver_debug("no packet data read, now sleeping\n");
			
			/* 等待唤醒 */
			if (wait_event_interruptible(packet_info->wait_queue, (packet_info->packet_cnt != 0)))
				return -ERESTARTSYS;

			test_driver_debug("wake up\n");

			mutex_lock(&packet_info->packet_mutex);		
		}
		/* 取出第一个节点并删除 */
		if (!list_empty(&(packet_info->packet.list))) {
			packet = list_first_entry(&(packet_info->packet.list), struct test_char_drv_packet, list);
			list_del(&packet->list);
		}

		if (packet_info->packet_cnt)
			packet_info->packet_cnt--;
		mutex_unlock(&packet_info->packet_mutex);		

		/* 唤醒其他等待线程 */
		wake_up_interruptible(&packet_info->wait_queue);

		/* 拷贝数据 */
		if (!packet) {
			test_driver_debug("list and len not match\n");
			return -EFAULT;
		}
		
		/* 取最小值，防止溢出 */
		count = min(nbytes, (size_t)(packet->len));
		if (copy_to_user(buf, packet->data, count)) {
			return -EFAULT;
		}
		vfree(packet->data);
		kfree(packet);
		
		return count;
		
		break;	

	case TEST_CHAR_NORMAL_MODE:
	case TEST_CHAR_PIPE_MODE:
		if (priv->mode == TEST_CHAR_PIPE_MODE) {
			buf_info = radix_tree_lookup(&dev->buf_tree, (unsigned long)file);
			if (!buf_info) {
				test_driver_debug("radix_tree_lookup error file=0x%x...\n", file);
				return -1;
			}
		} else {
			buf_info = &dev->buf_info;
		}
			
		test_driver_debug("reading nbytes=%ld, read_index=%d, write_index=%d\n", 
			nbytes, buf_info->read_index, buf_info->write_index);
		
		if (down_interruptible(&buf_info->sem)) 
			return -ERESTARTSYS;
		
		while (buf_info->write_index == buf_info->read_index
				&& buf_info->total_read_len == buf_info->total_write_len) {
			/* 没有数据读取 */
			up(&buf_info->sem);
			if (file->f_flags & O_NONBLOCK)
				return -EAGAIN;
		
			test_driver_debug("no data read, now sleeping\n");
		
			/* 等待唤醒 */
			if (wait_event_interruptible(buf_info->wait_queue, (buf_info->write_index != buf_info->read_index)))
				return -ERESTARTSYS;
		
			test_driver_debug("wake up\n");
		
			if (down_interruptible(&buf_info->sem)) 
				return -ERESTARTSYS;
		}
		
		/* 处理正常情况 */
		if (buf_info->write_index > buf_info->read_index) {
			nbytes = min(nbytes, (size_t)(buf_info->write_index - buf_info->read_index));	
			if (copy_to_user(buf, buf_info->buf_data + buf_info->read_index, nbytes)) {
				up(&buf_info->sem);
				return -EFAULT;
			}
		
			/* 设置读下标 */
			buf_info->read_index += nbytes;
		}
		else { /* 处理回头情况 */
			count = min(nbytes, (size_t)(buf_info->buf_size - buf_info->read_index));
			if (copy_to_user(buf, buf_info->buf_data + buf_info->read_index, count)) {
				up(&buf_info->sem);
				return -EFAULT;
			}
		
			buf_info->read_index += count;
		
			/* 读到了末尾 */
			if (buf_info->read_index >= buf_info->buf_size)
				buf_info->read_index = 0;
		
			/* 仍然需要读 */
			if (nbytes > count) {
				nbytes = min(nbytes - count, (size_t)(buf_info->write_index - buf_info->read_index));	
				if (copy_to_user(buf, buf_info->buf_data + buf_info->read_index, nbytes)) {
					up(&buf_info->sem);
					return -EFAULT;
				}
				buf_info->read_index += nbytes;
				count += nbytes;
				nbytes = count;
			}
		}
		
		/* 读到了末尾 */
		if (buf_info->read_index >= buf_info->buf_size)
			buf_info->read_index = 0;
		
		buf_info->total_read_len += nbytes;
		
		up(&buf_info->sem);
		
		test_driver_debug("readed nbytes=%ld, read_index=%d, write_index=%d\n", 
			nbytes, buf_info->read_index, buf_info->write_index);
		
		wake_up_interruptible(&buf_info->wait_queue);
			
		return nbytes;
		
		break;
	default:
		break;
	}

	return -1;
}

static ssize_t test_char_drv_write(struct file *file, const char __user *buf,
						 size_t nbytes, loff_t *ppos)
{
	size_t count;
	size_t total_write_count = 0;
	struct test_char_priv *priv = file->private_data;
	struct test_char_drv_dev *dev = priv->dev;
	struct test_char_drv_buf_info *buf_info; 
	struct test_char_drv_packet_info *packet_info = NULL; 
	struct test_char_drv_packet *packet; 
	unsigned int packet_cnt;

	test_driver_debug("write file=0x%x pid=%d, mode=%d\n", file, current->pid, priv->mode);

	switch (priv->mode)
	{
	case TEST_CHAR_PACKET_MODE:
	case TEST_CHAR_ROUND_ROBIN_PACKET_MODE:	
		if (priv->mode == TEST_CHAR_PACKET_MODE) 
			packet_info = &dev->packet_info;
		else {
			struct list_head *list = NULL;
			unsigned int index = 0;

			/* round robin  */
			mutex_lock(&dev->round_head_list_mutex);
			if(!list_empty(&dev->round_head_list)) {
				list_for_each(list, &dev->round_head_list) {
					index++;
					if (index > dev->round_index) {
						struct test_char_priv *priv_tmp;
						priv_tmp = list_entry(list, struct test_char_priv, list);
						packet_info = &priv_tmp->packet_info;
						dev->round_index++;
						break; /* 跳出循环 */
					}
				}
			}

			/* 回头 */
			if (dev->round_index >= dev->round_cnt || index == 0) {
				dev->round_index = 0;
			}
			mutex_unlock(&dev->round_head_list_mutex);
		}

		if (!packet_info) {
			test_driver_debug("no recv packet_info error ...\n");
			return -1;
		}
		
		test_driver_debug("now copy packet data packet_info=0x%x...\n", packet_info);
		
		/* 限制个数 */
		mutex_lock(&packet_info->packet_mutex);
		packet_cnt = packet_info->packet_cnt;
 		mutex_unlock(&packet_info->packet_mutex);
		if(packet_cnt > 100) {
			test_driver_debug("packet_cnt more than 100 error ...\n");
			return -1;
		}

		packet = kmalloc(sizeof(struct test_char_drv_packet), GFP_KERNEL);
		if (!packet) {
			test_driver_debug("kmalloc test_char_drv_packet error ...\n");
			return -1;
		}

		packet->data = vmalloc(nbytes);
		if (!packet->data) {
			kfree(packet);
			test_driver_debug("kmalloc nbytes error ...\n");
			return -1;
		}
		
		if (copy_from_user(packet->data, buf, nbytes)) {
			vfree(packet->data);
			kfree(packet);
			test_driver_debug("copy_from_user error ...\n");
			return -1;
		}

		packet->len = nbytes;
		INIT_LIST_HEAD(&packet->list);
		
		/* 加入链表中 */
		mutex_lock(&packet_info->packet_mutex);
		list_add_tail(&packet->list, &(packet_info->packet.list));
		packet_info->packet_cnt++;
 		mutex_unlock(&packet_info->packet_mutex);

		/* 唤醒正在等待的线程 */
		wake_up_interruptible(&packet_info->wait_queue);

		return nbytes;
		
		break;
	case TEST_CHAR_NORMAL_MODE:
	case TEST_CHAR_PIPE_MODE:
		if (priv->mode == TEST_CHAR_PIPE_MODE) {
			buf_info = radix_tree_lookup(&dev->buf_tree, (unsigned long)file);
			if (!buf_info) {
				test_driver_debug("radix_tree_lookup error file=0x%x...\n", file);
				return -1;
			}
		} else {
			buf_info = &dev->buf_info;
		}
		
		test_driver_debug("writing nbytes=%ld, read_index=%d, write_index=%d\n", 
			nbytes, buf_info->read_index, buf_info->write_index);

		if (down_interruptible(&buf_info->sem)) 
			return -ERESTARTSYS;

	copy_loop:
		while (buf_info->write_index == buf_info->read_index 
				&& buf_info->total_read_len != buf_info->total_write_len) {
			up(&buf_info->sem);
			if (file->f_flags & O_NONBLOCK)
				return -EAGAIN;

			test_driver_debug("no space write, now sleeping\n");

			/* 等待唤醒 */
			if (wait_event_interruptible(buf_info->wait_queue, (buf_info->write_index != buf_info->read_index)))
				return -ERESTARTSYS;

			test_driver_debug("wake up\n");

			if (down_interruptible(&buf_info->sem)) 
				return -ERESTARTSYS;
		}
				
		if (buf_info->write_index >= buf_info->read_index) {
			count = min(nbytes, (size_t)(buf_info->buf_size - buf_info->write_index));
		}
		else {
			count = min(nbytes, (size_t)(buf_info->read_index - buf_info->write_index));
		}
		
		if (copy_from_user(buf_info->buf_data + buf_info->write_index, buf + total_write_count, count)) {
			up(&buf_info->sem);
			return -EFAULT;
		}
		nbytes -= count;
		buf_info->write_index += count;
		buf_info->total_write_len += count;
		total_write_count += count;

		if (buf_info->write_index >= buf_info->buf_size) {
			buf_info->write_index = 0;
		}

		if(nbytes > 0)
			goto copy_loop;
		
		up(&buf_info->sem);

		test_driver_debug("writed nbytes=%ld, read_index=%d, write_index=%d\n", 
			nbytes, buf_info->read_index, buf_info->write_index);
		
		wake_up_interruptible(&buf_info->wait_queue);

		return total_write_count;		
	
		break;

	default:
		break;
	}

	return -1;	
}


#define TEST_CHAR_DRV_MAGIC                          'k'
#define TEST_CHAR_DRV_CMD_INIT_PIPE                  _IO(TEST_CHAR_DRV_MAGIC, 0x1a)
#define TEST_CHAR_DRV_CMD_INIT_PACKET                _IO(TEST_CHAR_DRV_MAGIC, 0x1b)
#define TEST_CHAR_DRV_CMD_INIT_ROUND_ROBIN_WRITE_PACKET    _IO(TEST_CHAR_DRV_MAGIC, 0x1c)
#define TEST_CHAR_DRV_CMD_INIT_ROUND_ROBIN_READ_PACKET    _IO(TEST_CHAR_DRV_MAGIC, 0x1d)

long test_char_drv_unlocked_ioctl(struct file *file, unsigned int cmd, unsigned long arg)
{
	struct test_char_priv *priv;
	struct test_char_drv_dev *dev;
	struct test_char_drv_buf_info *buf_info; 
	struct test_char_drv_packet_info *packet_info;

	test_driver_debug("ioctl begin cmd=0x%x, TEST_CHAR_DRV_CMD_INIT_PIPE=0x%x\n", 
			cmd, TEST_CHAR_DRV_CMD_INIT_PIPE);
#if 1
	switch (cmd) {
	case TEST_CHAR_DRV_CMD_INIT_PIPE:

		/* 父子进程file是同一个 */
		priv = file->private_data;		
		dev  = priv->dev;

		test_driver_debug("ioctl begin file=0x%x count_in_buf_tree=%d\n", file, atomic64_read(&dev->count_in_buf_tree));
		
		/* 查找是否存在 */
		buf_info = radix_tree_lookup(&dev->buf_tree, (unsigned long)file);
		if (!buf_info && atomic64_read(&dev->count_in_buf_tree) < 100) {
			buf_info = kmalloc(sizeof(struct test_char_drv_buf_info), GFP_KERNEL);
			if (buf_info) {
				unsigned long flags;
				struct test_char_drv_buf_info *tmp;

				test_driver_debug("set pipe buf_info begin \n");
				
				/* 设置变量 */
				memset(buf_info, 0, sizeof(struct test_char_drv_buf_info));
				buf_info->buf_size = dev->buf_info.buf_size;
				buf_info->buf_data = vmalloc(buf_info->buf_size);
				if (!buf_info->buf_data) {
					kfree(buf_info);
					return -1;
				}
				sema_init(&buf_info->sem, 1);
				init_waitqueue_head(&buf_info->wait_queue);

				/* 插入到表中 */
				spin_lock_irqsave(&dev->tree_lock, flags);
				tmp = radix_tree_lookup(&dev->buf_tree, (unsigned long)file);
				if (tmp) {
					/* 已经存在 */
					vfree(buf_info->buf_data);
					kfree(buf_info);
				} else {
					radix_tree_insert(&dev->buf_tree, (unsigned long)file, buf_info);
					atomic64_inc(&dev->count_in_buf_tree);
					/* 切换到PIPE模式 */
					priv->mode = TEST_CHAR_PIPE_MODE;
				}
				spin_unlock_irqrestore(&dev->tree_lock, flags);

				test_driver_debug("set pipe buf_info end \n");
		   }
		}
			
		break;
	case TEST_CHAR_DRV_CMD_INIT_PACKET:
		priv = file->private_data;		
		priv->mode = TEST_CHAR_PACKET_MODE;
		break;
	case TEST_CHAR_DRV_CMD_INIT_ROUND_ROBIN_WRITE_PACKET:
		priv = file->private_data;			
		/**/
		priv->mode = TEST_CHAR_ROUND_ROBIN_PACKET_MODE;
		break;
	case TEST_CHAR_DRV_CMD_INIT_ROUND_ROBIN_READ_PACKET:
		priv = file->private_data;			
		dev  = priv->dev;
		/**/
		priv->mode = TEST_CHAR_ROUND_ROBIN_PACKET_MODE;

		INIT_LIST_HEAD(&(priv->list));
		
		/* 初始化packet */
		packet_info = &priv->packet_info;
		packet_info->packet_cnt = 0;
		INIT_LIST_HEAD(&(packet_info->packet.list));
		mutex_init(&packet_info->packet_mutex);
		init_waitqueue_head(&packet_info->wait_queue);

		/* 加入到全局链表中 */
		mutex_lock(&dev->round_head_list_mutex);
		list_add_tail(&(priv->list), &dev->round_head_list);
		dev->round_cnt++;
		mutex_unlock(&dev->round_head_list_mutex);
		
		break;	
	default:
		break;
	}
#endif

	return 0;
}
						 

static struct file_operations test_char_drv_fops = {
    .open = test_char_drv_open,
    .read = test_char_drv_read,
    .write = test_char_drv_write,
    .unlocked_ioctl = test_char_drv_unlocked_ioctl,
    .release = test_char_drv_release,
};
	
int test_char_drv_init(void)
{
	int err;
	struct test_char_drv_buf_info *buf_info;
	struct test_char_drv_packet_info *packet_info;
		
	test_driver_debug("int begin...\n");

	memset(&char_dev, 0, sizeof(char_dev));

	buf_info = &(char_dev.buf_info);

	buf_info->buf_size = 4*1024*1024;
	buf_info->read_index = 0;
	buf_info->write_index = 0;
	buf_info->total_read_len = 0;
	buf_info->total_write_len = 0;
	
	/* 先申请内存 */
	buf_info->buf_data = vmalloc(buf_info->buf_size);
	if (!buf_info->buf_data) {
		test_driver_debug("vmalloc buf_data size=%ld error!\n", buf_info->buf_size);
		return -1;
	}
	/* 设置默认值 */
	memset(buf_info->buf_data, 0, buf_info->buf_size);

	/* packet info 初始化 */
	packet_info = &(char_dev.packet_info);
	packet_info->packet_cnt = 0;
	INIT_LIST_HEAD(&(packet_info->packet.list));
	mutex_init(&packet_info->packet_mutex);
	init_waitqueue_head(&packet_info->wait_queue);

	INIT_LIST_HEAD(&(char_dev.round_head_list));
	mutex_init(&(char_dev.round_head_list_mutex));
	char_dev.round_index = 0;
	char_dev.round_cnt   = 0;
	
    /* 注册字符设备 */
    if (register_chrdev(major_id, "test_char_drv_03", &test_char_drv_fops)) {
		test_driver_debug("unable to get major %d!", major_id);
		goto out_free;
    }	

	INIT_RADIX_TREE(&(char_dev.buf_tree), GFP_ATOMIC);
	spin_lock_init(&(char_dev.tree_lock));
	atomic64_set(&(char_dev.count_in_buf_tree), 0);

	sema_init(&buf_info->sem, 1);
	init_waitqueue_head(&buf_info->wait_queue);

	test_driver_debug("init success...\n");

    return 0;

out_free:
	if(buf_info->buf_data) 
		vfree(buf_info->buf_data);
	
	return -1;
}

void test_char_drv_exit(void)
{
	struct test_char_drv_buf_info *buf_info;

	buf_info = &(char_dev.buf_info);
	
	test_driver_debug("exit begin...\n");
	
	if(buf_info->buf_data) 
		vfree(buf_info->buf_data);

	/* 释放其他内存区域 */

	/* 释放radix_tree */

    unregister_chrdev(major_id, "test_char_drv_03");

	test_driver_debug("exit success...\n");
}

module_init(test_char_drv_init);
module_exit(test_char_drv_exit);
MODULE_AUTHOR("yankun");
MODULE_LICENSE("GPL");

